# -*- coding: utf-8 -*-
"""
彗星亮度预测工具 v7.1
- 删除 3D 轨道窗口
- 新增计算进度条 + 线程化计算
- 修复 name_entry 初始化顺序 bug
"""
import threading
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from tkinter import ttk, messagebox
import tkinter as tk
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
matplotlib.use('TkAgg')

# ---------- 中文显示 ----------
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False


# ---------- 主程序 ----------
class CometBrightnessApp:
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("彗星亮度预测工具 v7.1")
        self.root.geometry("1400x900")
        self.root.minsize(1200, 700)

        # 数据缓存
        self.current_dates = None
        self.cobs_calc_m_pred = None
        self.cobs_obs_m_pred = None
        self.astro_m_pred = None
        self.mpc_m_pred = None
        self.r_pred = None
        self.delta_pred = None

        # 子窗口引用
        self.brightness_window = None
        self.distance_window = None

        # 进度条变量
        self.progress_var = tk.IntVar(value=0)
        self.progress_text = tk.StringVar(value="就绪")

        # 构建 UI
        self._build_ui()

        # 延迟加载默认数据，确保 UI 已构建完成
        self.root.after(100, self.fill_sample_data)

    # -------------------- UI 构建 --------------------
    def _build_ui(self):
        main_paned = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        main_paned.pack(fill=tk.BOTH, expand=True)

        left_frame = ttk.Frame(main_paned, width=450)
        main_paned.add(left_frame, weight=0)

        right_frame = ttk.Frame(main_paned)
        main_paned.add(right_frame, weight=1)

        # 左侧：参数 + 控制
        self._build_left(left_frame)

        # 右侧：结果 + 进度
        self._build_right(right_frame)

    def _build_left(self, master):
        notebook = ttk.Notebook(master)
        notebook.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        basic = ttk.Frame(notebook)
        notebook.add(basic, text="基本参数")
        self._build_basic_page(basic)

        model = ttk.Frame(notebook)
        notebook.add(model, text="模型参数")
        self._build_model_page(model)

        ctrl = ttk.Frame(notebook)
        notebook.add(ctrl, text="预测控制")
        self._build_control_page(ctrl)

    def _build_basic_page(self, master):
        # 彗星选择
        comet_frame = ttk.LabelFrame(master, text="彗星选择", padding=10)
        comet_frame.pack(fill=tk.X, pady=5)
        ttk.Label(comet_frame, text="选择彗星:").grid(
            row=0, column=0, sticky=tk.W, padx=5)
        self.comet_var = tk.StringVar()
        cb = ttk.Combobox(
            comet_frame, textvariable=self.comet_var, width=25, state="readonly")
        cb["values"] = [
            "323P-B/SOHO", "240P/NEAT", "24P/Schaumasse",
            "C/2025 A6 (Lemmon)", "C/2025 R2 (SWAN)", "C/2025 K1 (ATLAS)",
            "210P/Christensen", "3I/ATLAS", "C/2024 E1 (Wierzchos)"
        ]
        cb.grid(row=0, column=1, padx=5)
        cb.bind("<<ComboboxSelected>>", self.on_comet_selected)

        # 彗星信息
        info_frame = ttk.LabelFrame(master, text="彗星信息", padding=10)
        info_frame.pack(fill=tk.X, pady=5)
        ttk.Label(info_frame, text="彗星名称:").grid(
            row=0, column=0, sticky=tk.W, padx=5)
        self.name_entry = ttk.Entry(info_frame, width=25)
        self.name_entry.grid(row=0, column=1, padx=5)

        # 轨道元素
        orbit_frame = ttk.LabelFrame(master, text="轨道元素 (MPC格式)", padding=10)
        orbit_frame.pack(fill=tk.X, pady=5)
        labels = [
            ("q (近日点距离 AU)", "perihelion_distance"),
            ("a (半长轴 AU)", "semi_major_axis"),
            ("e (偏心率)", "eccentricity"),
            ("ω (近日点幅角 °)", "arg_perihelion"),
            ("Ω (升交点赤经 °)", "long_asc_node"),
            ("i (轨道倾角 °)", "inclination"),
            ("P (轨道周期 年)", "orbital_period"),
            ("T (近日点时间)", "perihelion_time"),
            ("历元时间", "epoch_time")
        ]
        self.orbit_entries = {}
        for idx, (txt, key) in enumerate(labels):
            ttk.Label(orbit_frame, text=txt).grid(
                row=idx, column=0, sticky=tk.W, padx=5)
            ent = ttk.Entry(orbit_frame, width=18)
            ent.grid(row=idx, column=1, padx=5)
            self.orbit_entries[key] = ent

    def _build_model_page(self, master):
        switch = ttk.LabelFrame(master, text="模型开关", padding=10)
        switch.pack(fill=tk.X, pady=5)
        self.use_cobs_calc_var = tk.BooleanVar(value=True)
        self.use_cobs_obs_var = tk.BooleanVar(value=True)
        self.use_astro_var = tk.BooleanVar(value=True)
        self.use_mpc_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(switch, text="启用 MPC 模型",
                        variable=self.use_mpc_var).pack(anchor=tk.W)
        ttk.Checkbutton(switch, text="启用 COBS-Calc 模型",
                        variable=self.use_cobs_calc_var).pack(anchor=tk.W)
        ttk.Checkbutton(switch, text="启用 COBS-Observe 模型",
                        variable=self.use_cobs_obs_var).pack(anchor=tk.W)
        ttk.Checkbutton(switch, text="启用 ASTRO 模型",
                        variable=self.use_astro_var).pack(anchor=tk.W)

        mag = ttk.LabelFrame(master, text="星等参数 (H/n)", padding=10)
        mag.pack(fill=tk.X, pady=5)
        self.cobs_calc_H_entry = self._add_hn_row(mag, 0, "COBS-Calc")
        self.cobs_obs_H_entry = self._add_hn_row(mag, 1, "COBS-Observe")
        self.astro_H_entry = self._add_hn_row(mag, 2, "ASTRO")
        self.mpc_H_entry = self._add_hn_row(mag, 3, "MPC")

    def _add_hn_row(self, master, row, label):
        ttk.Label(master, text=f"{label} H:").grid(
            row=row, column=0, padx=5, sticky=tk.W)
        h = ttk.Entry(master, width=8)
        h.grid(row=row, column=1, padx=5)
        ttk.Label(master, text="n:").grid(row=row, column=2, padx=5)
        n = ttk.Entry(master, width=8)
        n.grid(row=row, column=3, padx=5)
        return h, n

    def _build_control_page(self, master):
        mode = ttk.LabelFrame(master, text="预测模式", padding=10)
        mode.pack(fill=tk.X, pady=5)
        self.pred_mode = tk.StringVar(value="days")
        ttk.Radiobutton(mode, text="按天数预测", variable=self.pred_mode,
                        value="days").pack(anchor=tk.W)
        ttk.Radiobutton(mode, text="按日期范围预测", variable=self.pred_mode,
                        value="dates").pack(anchor=tk.W)
        self.pred_mode.trace_add("write", self.toggle_pred_mode)

        self.days_frame = ttk.Frame(mode)
        self.days_frame.pack(fill=tk.X, pady=5)
        ttk.Label(self.days_frame, text="预测天数:").pack(side=tk.LEFT, padx=5)
        self.pred_days_entry = ttk.Entry(self.days_frame, width=10)
        self.pred_days_entry.insert(0, "365")
        self.pred_days_entry.pack(side=tk.LEFT, padx=5)

        self.dates_frame = ttk.Frame(mode)
        ttk.Label(self.dates_frame, text="开始日期:").pack(side=tk.LEFT, padx=5)
        self.start_date_entry = ttk.Entry(self.dates_frame, width=12)
        self.start_date_entry.pack(side=tk.LEFT, padx=5)
        ttk.Label(self.dates_frame, text="结束日期:").pack(side=tk.LEFT, padx=5)
        self.end_date_entry = ttk.Entry(self.dates_frame, width=12)
        self.end_date_entry.pack(side=tk.LEFT, padx=5)

        btn = ttk.Frame(master)
        btn.pack(fill=tk.X, pady=10)
        ttk.Button(btn, text="开始预测", command=self.run_prediction).pack(
            side=tk.LEFT, padx=5)
        ttk.Button(btn, text="重置参数", command=self.clear_entries).pack(
            side=tk.LEFT, padx=5)

    def _build_right(self, master):
        top = ttk.Frame(master)
        top.pack(fill=tk.BOTH, expand=True)

        self._build_result(top)

        bottom = ttk.Frame(master, height=40)
        bottom.pack(fill=tk.X, side=tk.BOTTOM, padx=5, pady=2)
        self.progress_bar = ttk.Progressbar(
            bottom, variable=self.progress_var, maximum=100)
        self.progress_bar.pack(fill=tk.X, side=tk.LEFT,
                               expand=True, padx=(0, 8))
        ttk.Label(bottom, textvariable=self.progress_text,
                  width=20).pack(side=tk.RIGHT)

        chart = ttk.LabelFrame(top, text="图表窗口", padding=10)
        chart.pack(fill=tk.X, pady=5)
        for txt, cmd in [("打开亮度图表", self.open_brightness_window),
                         ("打开距离图表", self.open_distance_window),
                         ("关闭所有图表", self.close_all_windows)]:
            ttk.Button(chart, text=txt, command=cmd).pack(fill=tk.X, pady=3)

    def _build_result(self, master):
        result = ttk.LabelFrame(master, text="预测结果", padding=10)
        result.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        canvas = tk.Canvas(result)
        scroll = ttk.Scrollbar(result, orient="vertical", command=canvas.yview)
        scrollable = ttk.Frame(canvas)
        scrollable.bind("<Configure>", lambda e: canvas.configure(
            scrollregion=canvas.bbox("all")))
        canvas.create_window((0, 0), window=scrollable, anchor="nw")
        canvas.configure(yscrollcommand=scroll.set)
        canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scroll.pack(side=tk.RIGHT, fill=tk.Y)

        items = [
            "轨道类型", "近日点日期", "预测总天数", "近地点距离",
            "COBS-Calc最亮星等", "COBS-Calc最亮日期",
            "COBS-Observe最亮星等", "COBS-Observe最亮日期",
            "ASTRO最亮星等", "ASTRO最亮日期",
            "MPC最亮星等", "MPC最亮日期",
            "今日日彗距", "今日地彗距",
            "今日COBS-Calc星等", "今日COBS-Observe星等", "今日ASTRO星等", "今日MPC星等",
            "30天后日彗距", "30天后地彗距", "30天后COBS-Calc星等", "30天后COBS-Observe星等", "30天后ASTRO星等", "30天后MPC星等",
            "60天后日彗距", "60天后地彗距", "60天后COBS-Calc星等", "60天后COBS-Observe星等", "60天后ASTRO星等", "60天后MPC星等",
            "彗发长度 (km)", "彗发长度 (°)", "彗尾长度 (km)", "彗尾长度 (°)"
        ]
        self.info_labels = {}
        for k in items:
            f = ttk.Frame(scrollable)
            f.pack(fill=tk.X, pady=1)
            ttk.Label(f, text=k + ":", width=26,
                      anchor="e").pack(side=tk.LEFT, padx=5)
            lab = ttk.Label(f, text="待计算", width=40, anchor="w")
            lab.pack(side=tk.LEFT, padx=5)
            self.info_labels[k] = lab

    # -------------------- 事件 --------------------
    def toggle_pred_mode(self, *_):
        if self.pred_mode.get() == "days":
            self.days_frame.pack(fill=tk.X, pady=5)
            self.dates_frame.pack_forget()
        else:
            self.days_frame.pack_forget()
            self.dates_frame.pack(fill=tk.X, pady=5)
            if not self.start_date_entry.get():
                self.start_date_entry.insert(
                    0, datetime.now().strftime('%Y-%m-%d'))
            if not self.end_date_entry.get():
                self.end_date_entry.insert(
                    0, (datetime.now() + timedelta(days=1825)).strftime('%Y-%m-%d'))

    def on_comet_selected(self, *_):
        self.fill_comet_data(self.comet_var.get())

    # -------------------- 数据 --------------------
    def fill_sample_data(self):
        self.comet_var.set("323P-B/SOHO")
        self.fill_comet_data("323P-B/SOHO")

    def fill_comet_data(self, name):
        data = COMET_SAMPLE_DATA.get(name)
        if not data:
            return
        self.clear_entries()
        self.name_entry.insert(0, data["name"])
        for k, v in data["orbit"].items():
            if v:
                self.orbit_entries[k].insert(0, v)
        self.use_cobs_calc_var.set(data["model_params"]["use_cobs_calc"])
        self.use_cobs_obs_var.set(data["model_params"]["use_cobs_obs"])
        self.use_astro_var.set(data["model_params"]["use_astro"])
        self.use_mpc_var.set(data["model_params"]["use_mpc"])
        self.cobs_calc_H_entry[0].insert(
            0, data["model_params"]["cobs_calc_H"])
        self.cobs_calc_H_entry[1].insert(
            0, data["model_params"]["cobs_calc_n"])
        self.cobs_obs_H_entry[0].insert(0, data["model_params"]["cobs_obs_H"])
        self.cobs_obs_H_entry[1].insert(0, data["model_params"]["cobs_obs_n"])
        self.astro_H_entry[0].insert(0, data["model_params"]["astro_H"])
        self.astro_H_entry[1].insert(0, data["model_params"]["astro_n"])
        self.mpc_H_entry[0].insert(0, data["model_params"]["mpc_H"])
        self.mpc_H_entry[1].insert(0, data["model_params"]["mpc_n"])

    def clear_entries(self):
        for e in self.orbit_entries.values():
            e.delete(0, tk.END)
        for h, n in [self.cobs_calc_H_entry, self.cobs_obs_H_entry, self.astro_H_entry, self.mpc_H_entry]:
            h.delete(0, tk.END)
            n.delete(0, tk.END)
        for lab in self.info_labels.values():
            lab.config(text="待计算")

    # -------------------- 校验 --------------------
    def validate_inputs(self):
        try:
            name = self.name_entry.get().strip()
            if not name:
                raise ValueError("请输入彗星名称")
            q = float(self.orbit_entries["perihelion_distance"].get())
            e = float(self.orbit_entries["eccentricity"].get())
            peri_str = self.orbit_entries["perihelion_time"].get()
            a_str = self.orbit_entries["semi_major_axis"].get()
            a = float(a_str) if a_str else self.calc_semi_major_axis(q, e)
            peri_date = datetime.strptime(peri_str, "%Y-%m-%d %H:%M:%S")
            calc_H, calc_n = self.cobs_calc_H_entry[0].get(
            ), self.cobs_calc_H_entry[1].get()
            obs_H, obs_n = self.cobs_obs_H_entry[0].get(
            ), self.cobs_obs_H_entry[1].get()
            astro_H, astro_n = self.astro_H_entry[0].get(
            ), self.astro_H_entry[1].get()
            mpc_H, mpc_n = self.mpc_H_entry[0].get(), self.mpc_H_entry[1].get()
            if not 0.001 <= q <= 1000:
                raise ValueError("近日点距离超出 0.001-1000 AU")
            if e < 0:
                raise ValueError("偏心率不能为负")
            if self.use_cobs_calc_var.get() and calc_H and calc_n:
                if not (-50 <= float(calc_H) <= 5000 and -50 <= float(calc_n) <= 5000):
                    raise ValueError("COBS-Calc 参数超出 -50~5000")
            if self.use_cobs_obs_var.get() and obs_H and obs_n:
                if not (-50 <= float(obs_H) <= 5000 and -50 <= float(obs_n) <= 5000):
                    raise ValueError("COBS-Observe 参数超出 -50~5000")
            if self.use_astro_var.get() and astro_H and astro_n:
                if not (-50 <= float(astro_H) <= 5000 and -50 <= float(astro_n) <= 5000):
                    raise ValueError("ASTRO 参数超出 -50~5000")
            if self.use_mpc_var.get() and mpc_H and mpc_n:
                if not (-50 <= float(mpc_H) <= 5000 and -50 <= float(mpc_n) <= 5000):
                    raise ValueError("MPC 参数超出 -50~5000")
            return (name, q, a, e, peri_date, calc_H, calc_n, obs_H, obs_n, astro_H, astro_n, mpc_H, mpc_n)
        except Exception as e:
            messagebox.showerror("输入错误", str(e))
            return None

    def calc_semi_major_axis(self, q, e):
        if e < 1:
            return q / (1 - e)
        if e == 1:
            return 1e10
        return q / (1 - e)

    # -------------------- 预测 --------------------
    def run_prediction(self):
        inputs = self.validate_inputs()
        if inputs is None:
            return
        threading.Thread(target=self._do_prediction,
                         args=(inputs,), daemon=True).start()

    def _do_prediction(self, inputs):
        self._set_progress(0, "开始计算……")
        try:
            (name, q, a, e, peri_date, calc_H, calc_n, obs_H,
             obs_n, astro_H, astro_n, mpc_H, mpc_n) = inputs
            if self.pred_mode.get() == "days":
                future_days = int(self.pred_days_entry.get())
                if not 1 <= future_days <= 1_000_000:
                    raise ValueError("预测天数 1-1,000,000")
                start_date = peri_date - timedelta(days=future_days // 2)
                all_dates = [start_date +
                             timedelta(days=i) for i in range(future_days)]
            else:
                start_date = datetime.strptime(
                    self.start_date_entry.get(), '%Y-%m-%d')
                end_date = datetime.strptime(
                    self.end_date_entry.get(), '%Y-%m-%d')
                delta_days = (end_date - start_date).days
                if delta_days > 1_000_000:
                    raise ValueError("预测范围超过 1,000,000 天")
                all_dates = [start_date +
                             timedelta(days=i) for i in range(delta_days + 1)]
            total = len(all_dates)
            self._set_progress(5, "计算日彗距/地彗距……")

            r_pred = []
            for i, date in enumerate(all_dates):
                days_from_peri = (date - peri_date).days
                if e < 1:
                    M = 2 * np.pi * days_from_peri / (a ** 1.5 * 365.25)
                    E = self.solve_kepler(M, e)
                    r = a * (1 - e * np.cos(E))
                elif e == 1:
                    r = self.calc_parabolic(days_from_peri, q)
                else:
                    r = self.calc_hyperbolic(days_from_peri, q, e, a)
                r_pred.append(max(r, 0.001))
                if i % 100 == 0:
                    self._set_progress(5 + 25 * i // total, "计算日彗距/地彗距……")
            r_pred = np.array(r_pred)

            delta_pred = []
            for i, date in enumerate(all_dates):
                delta = self.calc_earth_distance(date, r_pred[i], e, peri_date)
                delta_pred.append(delta)
                if i % 100 == 0:
                    self._set_progress(30 + 20 * i // total, "计算地彗距……")
            delta_pred = np.array(delta_pred)

            self._set_progress(55, "计算星等……")
            self.cobs_calc_m_pred = None
            self.cobs_obs_m_pred = None
            self.astro_m_pred = None
            self.mpc_m_pred = None
            if self.use_cobs_calc_var.get() and calc_H and calc_n:
                self.cobs_calc_m_pred = float(
                    calc_H) + 5 * np.log10(delta_pred) + 2.5 * float(calc_n) * np.log10(r_pred)
                self.cobs_calc_m_pred = np.clip(
                    self.cobs_calc_m_pred, -100, 5000)
            if self.use_cobs_obs_var.get() and obs_H and obs_n:
                self.cobs_obs_m_pred = float(
                    obs_H) + 5 * np.log10(delta_pred) + 2.5 * float(obs_n) * np.log10(r_pred)
                self.cobs_obs_m_pred = np.clip(
                    self.cobs_obs_m_pred, -100, 5000)
            if self.use_astro_var.get() and astro_H and astro_n:
                self.astro_m_pred = float(
                    astro_H) + 5 * np.log10(delta_pred) + 2.5 * float(astro_n) * np.log10(r_pred)
                self.astro_m_pred = np.clip(self.astro_m_pred, -100, 5000)
            if self.use_mpc_var.get() and mpc_H and mpc_n:
                self.mpc_m_pred = float(
                    mpc_H) + 5 * np.log10(delta_pred) + 2.5 * float(mpc_n) * np.log10(r_pred)
                self.mpc_m_pred = np.clip(self.mpc_m_pred, -100, 5000)

            self.current_dates = all_dates
            self.r_pred = r_pred
            self.delta_pred = delta_pred
            self._set_progress(90, "更新界面……")
            self.update_results(name, e, all_dates, r_pred, total)
            self.update_brightness_window()
            self.update_distance_window()
            self._set_progress(100, "完成")
            messagebox.showinfo("预测完成", "彗星亮度预测计算完成！ (v7.1)")
        except Exception as e:
            self._set_progress(0, "错误")
            messagebox.showerror("预测错误", str(e))

    def _set_progress(self, val, text):
        self.progress_var.set(val)
        self.progress_text.set(text)
        self.root.update_idletasks()

    # -------------------- 轨道算法 --------------------
    def solve_kepler(self, M, e, tol=1e-8, max_iter=100):
        if abs(M) < 1e-6:
            return M
        E = M if e < 0.8 else (np.pi if M > 0 else -np.pi)
        for _ in range(max_iter):
            delta = (E - e * np.sin(E) - M) / (1 - e * np.cos(E))
            E -= delta
            if abs(E) > 10 * np.pi:
                E = np.sign(E) * 10 * np.pi
            if abs(delta) < tol:
                break
        return E

    def calc_parabolic(self, days_from_peri, q):
        W = 3 * days_from_peri * np.sqrt(0.5 / q ** 3) / 2
        s = W
        for _ in range(10):
            s_new = W - s ** 3 / 3
            if abs(s_new - s) < 1e-8:
                break
            s = s_new
        return max(q * (1 + s ** 2), 0.001)

    def calc_hyperbolic(self, days_from_peri, q, e, a):
        N = 2 * np.pi * days_from_peri / ((-a) ** 1.5 * 365.25)
        H = self.solve_hyperbolic(N, e)
        return max(a * (1 - e * np.cosh(H)), 0.001)

    def solve_hyperbolic(self, N, e, tol=1e-8, max_iter=100):
        if abs(N) < 0.1:
            H = N
        else:
            H = np.sign(N) * np.log(2 * abs(N) / e + 1.8)
        for _ in range(max_iter):
            sinhH, coshH = np.sinh(H), np.cosh(H)
            f = e * sinhH - H - N
            df = e * coshH - 1
            if abs(df) < 1e-12:
                H += 0.1
                continue
            delta = f / df
            H -= delta
            if abs(H) > 100:
                H = np.sign(H) * 100
            if abs(delta) < tol:
                break
        return H

    def calc_earth_distance(self, date, r, e, peri_date):
        days_from_peri = (date - peri_date).days
        earth_angle = 2 * np.pi * (days_from_peri % 365.25) / 365.25
        earth_x, earth_y = np.cos(earth_angle), np.sin(earth_angle)
        omega = np.radians(float(self.orbit_entries["arg_perihelion"].get()))
        i = np.radians(float(self.orbit_entries["inclination"].get()))
        comet_x = r * np.cos(omega)
        comet_y = r * np.sin(omega) * np.cos(i)
        return max(np.sqrt((comet_x - earth_x) ** 2 + (comet_y - earth_y) ** 2), 0.001)

    # -------------------- 结果更新 --------------------
    def update_results(self, name, e, dates, r_pred, total):
        orbit_type = "椭圆" if e < 1 else ("抛物线" if e == 1 else "双曲线")
        self.info_labels["轨道类型"].config(text=f"{orbit_type}轨道 (e={e:.6f})")
        self.info_labels["预测总天数"].config(text=f"{total:,} 天")
        peri_idx = np.argmin(r_pred)
        self.info_labels["近日点日期"].config(
            text=dates[peri_idx].strftime('%Y-%m-%d'))
        self.info_labels["近地点距离"].config(text=f"{r_pred[peri_idx]:.7f} AU")

        today = datetime.now().date()
        today_idx = next((i for i, d in enumerate(
            dates) if d.date() == today), None)
        if today_idx is not None:
            r, delta = r_pred[today_idx], self.delta_pred[today_idx]
            self.info_labels["今日日彗距"].config(text=f"{r:.7f} AU")
            self.info_labels["今日地彗距"].config(text=f"{delta:.7f} AU")
            for model, lbl in zip(["cobs_calc", "cobs_obs", "astro", "mpc"],
                                  ["今日COBS-Calc星等", "今日COBS-Observe星等", "今日ASTRO星等", "今日MPC星等"]):
                pred = getattr(self, f"{model}_m_pred")
                if pred is not None:
                    self.info_labels[lbl].config(text=f"{pred[today_idx]:.2f}")
            coma_km, coma_deg, tail_km, tail_deg = self.calc_coma_tail(
                r, delta)
            self.info_labels["彗发长度 (km)"].config(text=f"{coma_km:.0f}")
            self.info_labels["彗发长度 (°)"].config(text=f"{coma_deg:.4f}")
            self.info_labels["彗尾长度 (km)"].config(text=f"{tail_km:.0f}")
            self.info_labels["彗尾长度 (°)"].config(text=f"{tail_deg:.4f}")

        for days in [30, 60]:
            future = datetime.now() + timedelta(days=days)
            idx = next((i for i, d in enumerate(dates)
                       if d.date() == future.date()), None)
            if idx is None:
                continue
            r, delta = r_pred[idx], self.delta_pred[idx]
            self.info_labels[f"{days}天后日彗距"].config(text=f"{r:.7f} AU")
            self.info_labels[f"{days}天后地彗距"].config(text=f"{delta:.7f} AU")
            for model, lbl in zip(["cobs_calc", "cobs_obs", "astro", "mpc"],
                                  [f"{days}天后COBS-Calc星等", f"{days}天后COBS-Observe星等", f"{days}天后ASTRO星等", f"{days}天后MPC星等"]):
                pred = getattr(self, f"{model}_m_pred")
                if pred is not None:
                    self.info_labels[lbl].config(text=f"{pred[idx]:.2f}")

        for model, lbl_mag, lbl_date in zip(["cobs_calc", "cobs_obs", "astro", "mpc"],
                                            ["COBS-Calc最亮星等", "COBS-Observe最亮星等",
                                                "ASTRO最亮星等", "MPC最亮星等"],
                                            ["COBS-Calc最亮日期", "COBS-Observe最亮日期", "ASTRO最亮日期", "MPC最亮日期"]):
            pred = getattr(self, f"{model}_m_pred")
            if pred is not None:
                min_idx = np.argmin(pred)
                self.info_labels[lbl_mag].config(text=f"{pred[min_idx]:.2f}")
                self.info_labels[lbl_date].config(
                    text=dates[min_idx].strftime('%Y-%m-%d'))

    def calc_coma_tail(self, r, delta):
        coma_km = 100_000 / np.sqrt(r)
        tail_km = 1_000_000 / np.sqrt(r)
        delta_km = delta * 149_597_870.7
        coma_deg = np.degrees(coma_km / delta_km)
        tail_deg = np.degrees(tail_km / delta_km)
        return coma_km, coma_deg, tail_km, tail_deg

    # -------------------- 图表窗口 --------------------
    def open_brightness_window(self):
        if self.brightness_window is None or not self.brightness_window.winfo_exists():
            self.brightness_window = tk.Toplevel(self.root)
            self.brightness_window.title("亮度预测图表 v7.1")
            self.fig_b, self.ax_b = plt.subplots(figsize=(12, 8))
            self.canvas_b = FigureCanvasTkAgg(
                self.fig_b, master=self.brightness_window)
            self.canvas_b.get_tk_widget().pack(fill=tk.BOTH, expand=True)
            NavigationToolbar2Tk(
                self.canvas_b, self.brightness_window).update()
            if self.current_dates is not None:
                self.update_brightness_window()
        else:
            self.brightness_window.lift()

    def open_distance_window(self):
        if self.distance_window is None or not self.distance_window.winfo_exists():
            self.distance_window = tk.Toplevel(self.root)
            self.distance_window.title("距离变化图表 v7.1")
            self.fig_d, self.ax_d = plt.subplots(figsize=(12, 8))
            self.canvas_d = FigureCanvasTkAgg(
                self.fig_d, master=self.distance_window)
            self.canvas_d.get_tk_widget().pack(fill=tk.BOTH, expand=True)
            NavigationToolbar2Tk(self.canvas_d, self.distance_window).update()
            if self.current_dates is not None:
                self.update_distance_window()
        else:
            self.distance_window.lift()

    def close_all_windows(self):
        for w in [self.brightness_window, self.distance_window]:
            if w and w.winfo_exists():
                w.destroy()

    def update_brightness_window(self):
        if (self.brightness_window is None or not self.brightness_window.winfo_exists()
                or self.current_dates is None):
            return
        self.ax_b.clear()
        e = float(self.orbit_entries["eccentricity"].get())
        orbit_type = "椭圆" if e < 1 else ("抛物线" if e == 1 else "双曲线")
        colors = ['#d62728', '#2ca02c', '#000080', '#ff7f0e']
        labels = ['COBS-Calc', 'COBS-Observe', 'ASTRO', 'MPC']
        preds = [self.cobs_calc_m_pred, self.cobs_obs_m_pred,
                 self.astro_m_pred, self.mpc_m_pred]
        uses = [self.use_cobs_calc_var.get(), self.use_cobs_obs_var.get(),
                self.use_astro_var.get(), self.use_mpc_var.get()]
        for pred, col, lab, use in zip(preds, colors, labels, uses):
            if pred is not None and use:
                self.ax_b.plot(self.current_dates, pred,
                               color=col, label=lab, linewidth=2)
                min_mag = pred.min()
                min_idx = np.argmin(pred)
                self.ax_b.plot(self.current_dates[min_idx], min_mag, 'o', color=col, markersize=8,
                               label=f'{lab}最亮: {min_mag:.2f}等')
        today = datetime.now().date()
        today_idx = next((i for i, d in enumerate(
            self.current_dates) if d.date() == today), None)
        thirty = datetime.now() + timedelta(days=30)
        thirty_idx = next((i for i, d in enumerate(
            self.current_dates) if d.date() == thirty.date()), None)
        for pred, col, lab, use in zip(preds, colors, labels, uses):
            if pred is None or not use:
                continue
            if today_idx is not None:
                self.ax_b.plot(self.current_dates[today_idx], pred[today_idx], 's', color=col, markersize=6,
                               label=f'今日{lab}: {pred[today_idx]:.2f}等')
            if thirty_idx is not None:
                self.ax_b.plot(self.current_dates[thirty_idx], pred[thirty_idx], 'D', color=col, markersize=6,
                               label=f'30天后{lab}: {pred[thirty_idx]:.2f}等')
        self.ax_b.set_ylabel("视星等")
        self.ax_b.set_title(
            f"{self.comet_var.get()} 亮度预测 - {orbit_type} (v7.1)", fontsize=14)
        self.ax_b.grid(True, alpha=0.3)
        self.ax_b.legend()
        self.ax_b.invert_yaxis()
        self.fig_b.autofmt_xdate()
        self.fig_b.tight_layout()
        self.canvas_b.draw()

    def update_distance_window(self):
        if (self.distance_window is None or not self.distance_window.winfo_exists()
                or self.current_dates is None):
            return
        self.ax_d.clear()
        e = float(self.orbit_entries["eccentricity"].get())
        orbit_type = "椭圆" if e < 1 else ("抛物线" if e == 1 else "双曲线")
        self.ax_d.plot(self.current_dates, self.r_pred,
                       'orange', label='日彗距', linewidth=2)
        self.ax_d.plot(self.current_dates, self.delta_pred,
                       'blue', label='地彗距', linewidth=2)

        today_idx = next((i for i, d in enumerate(
            self.current_dates) if d.date() == datetime.now().date()), None)
        if today_idx is not None:
            r, delta = self.r_pred[today_idx], self.delta_pred[today_idx]
            self.ax_d.plot(self.current_dates[today_idx], r, 's', color='purple', markersize=6,
                           label=f'今日日彗距: {r:.7f} AU')
            self.ax_d.plot(self.current_dates[today_idx], delta, 's', color='green', markersize=6,
                           label=f'今日地彗距: {delta:.7f} AU')
        peri_date_str = self.orbit_entries["perihelion_time"].get()
        if peri_date_str:
            try:
                peri_date = datetime.strptime(
                    peri_date_str, "%Y-%m-%d %H:%M:%S").date()
                peri_idx = next((i for i, d in enumerate(
                    self.current_dates) if d.date() == peri_date), None)
                if peri_idx is not None:
                    self.ax_d.plot(self.current_dates[peri_idx], self.r_pred[peri_idx], 'D', color='orange', markersize=6,
                                   label=f'近日点: {self.r_pred[peri_idx]:.7f} AU')
                    self.ax_d.plot(self.current_dates[peri_idx], self.delta_pred[peri_idx], 'D', color='cyan', markersize=6,
                                   label=f'近日点地彗距: {self.delta_pred[peri_idx]:.7f} AU')
            except:
                pass
        self.ax_d.set_ylabel("距离 (AU)")
        self.ax_d.set_xlabel("日期")
        self.ax_d.set_title(
            f"{self.comet_var.get()} 距离变化 - {orbit_type} (v7.1)", fontsize=14)
        self.ax_d.grid(True, alpha=0.3)
        self.ax_d.legend()
        self.fig_d.autofmt_xdate()
        self.fig_d.tight_layout()
        self.canvas_d.draw()


# -------------------- 样本数据 --------------------
COMET_SAMPLE_DATA = {
    "323P-B/SOHO": {
        "name": "323P-B/SOHO",
        "orbit": {
            "perihelion_distance": "0.040063",
            "semi_major_axis": "2.890131",
            "eccentricity": "0.986138",
            "arg_perihelion": "353.9691",
            "long_asc_node": "323.4650",
            "inclination": "5.4627",
            "orbital_period": "4.9135",
            "perihelion_time": "2025-12-16 05:37:06",
            "epoch_time": "2025-10-06"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "20.0",
            "cobs_calc_n": "10.0",
            "cobs_obs_H": "19.5",
            "cobs_obs_n": "9.5",
            "astro_H": "18.5",
            "astro_n": "9.0",
            "mpc_H": "21.0",
            "mpc_n": "5.0"
        }
    },
    "240P/NEAT": {
        "name": "240P/NEAT",
        "orbit": {
            "perihelion_distance": "2.121663",
            "semi_major_axis": "3.860939",
            "eccentricity": "0.450480",
            "arg_perihelion": "352.0765",
            "long_asc_node": "74.9114",
            "inclination": "23.5370",
            "orbital_period": "7.5868",
            "perihelion_time": "2025-12-19 22:41:13",
            "epoch_time": "2025-10-06"
        },
        "model_params": {
            "use_cobs_calc": False,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": False,
            "cobs_calc_H": "",
            "cobs_calc_n": "",
            "cobs_obs_H": "11.0",
            "cobs_obs_n": "1.8",
            "astro_H": "10.5",
            "astro_n": "1.5",
            "mpc_H": "",
            "mpc_n": ""
        }
    },
    "24P/Schaumasse": {
        "name": "24P/Schaumasse",
        "orbit": {
            "perihelion_distance": "1.183913",
            "semi_major_axis": "4.058333",
            "eccentricity": "0.708276",
            "arg_perihelion": "58.4855",
            "long_asc_node": "78.2724",
            "inclination": "11.5023",
            "orbital_period": "8.1760",
            "perihelion_time": "2026-01-08 08:15:47",
            "epoch_time": "2025-10-06"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "6.5",
            "cobs_calc_n": "35.0",
            "cobs_obs_H": "6.5",
            "cobs_obs_n": "14.0",
            "astro_H": "6.5",
            "astro_n": "14.0",
            "mpc_H": "7.0",
            "mpc_n": "3.0"
        }
    },
    "C/2025 A6 (Lemmon)": {
        "name": "C/2025 A6 (Lemmon)",
        "orbit": {
            "perihelion_distance": "0.529889",
            "semi_major_axis": "121.283818",
            "eccentricity": "0.995631",
            "arg_perihelion": "132.9699",
            "long_asc_node": "108.0978",
            "inclination": "143.6635",
            "orbital_period": "1335.7395",
            "perihelion_time": "2025-11-08 12:53:59",
            "epoch_time": "2025-10-04"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "13.2",
            "cobs_calc_n": "10",
            "cobs_obs_H": "6.9",
            "cobs_obs_n": "3.9",
            "astro_H": "13.2",
            "astro_n": "10",
            "mpc_H": "14.0",
            "mpc_n": "2.0"
        }
    },
    "C/2025 R2 (SWAN)": {
        "name": "C/2025 R2 (SWAN)",
        "orbit": {
            "perihelion_distance": "0.504053",
            "semi_major_axis": "",
            "eccentricity": "0.993413",
            "arg_perihelion": "308.3194",
            "long_asc_node": "335.3460",
            "inclination": "4.4736",
            "orbital_period": "",
            "perihelion_time": "2025-09-12 19:55:55",
            "epoch_time": "2025-10-04"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "12.4",
            "cobs_calc_n": "10",
            "cobs_obs_H": "8.8",
            "cobs_obs_n": "2.0",
            "astro_H": "12.4",
            "astro_n": "10",
            "mpc_H": "15.0",
            "mpc_n": "1.0"
        }
    },
    "C/2025 K1 (ATLAS)": {
        "name": "C/2025 K1 (ATLAS)",
        "orbit": {
            "perihelion_distance": "0.334197",
            "semi_major_axis": "",
            "eccentricity": "1.000263",
            "arg_perihelion": "271.0256",
            "long_asc_node": "97.5572",
            "inclination": "147.8645",
            "orbital_period": "",
            "perihelion_time": "2025-10-08 10:31:34",
            "epoch_time": "2025-10-04"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "12.2",
            "cobs_calc_n": "10",
            "cobs_obs_H": "12.4",
            "cobs_obs_n": "3.0",
            "astro_H": "12.2",
            "astro_n": "10",
            "mpc_H": "16.0",
            "mpc_n": "0.5"
        }
    },
    "210P/Christensen": {
        "name": "210P/Christensen",
        "orbit": {
            "perihelion_distance": "0.524416",
            "semi_major_axis": "3.160446",
            "eccentricity": "0.834069",
            "arg_perihelion": "345.9494",
            "long_asc_node": "93.7975",
            "inclination": "10.2873",
            "orbital_period": "5.6188",
            "perihelion_time": "2025-11-22 17:25:43",
            "epoch_time": "2025-10-04"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "13.5",
            "cobs_calc_n": "10",
            "cobs_obs_H": "13.5",
            "cobs_obs_n": "4.0",
            "astro_H": "13.5",
            "astro_n": "10",
            "mpc_H": "17.0",
            "mpc_n": "1.0"
        }
    },
    "3I/ATLAS": {
        "name": "3I/ATLAS",
        "orbit": {
            "perihelion_distance": "1.356324",
            "semi_major_axis": "",
            "eccentricity": "6.138574",
            "arg_perihelion": "128.0125",
            "long_asc_node": "322.1574",
            "inclination": "175.1131",
            "orbital_period": "",
            "perihelion_time": "2025-10-29 11:35:39",
            "epoch_time": "2025-11-21"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "11.8",
            "cobs_calc_n": "10",
            "cobs_obs_H": "6.8",
            "cobs_obs_n": "6.0",
            "astro_H": "11.8",
            "astro_n": "10",
            "mpc_H": "18.0",
            "mpc_n": "2.0"
        }
    },
    "C/2024 E1 (Wierzchos)": {
        "name": "C/2024 E1 (Wierzchos)",
        "orbit": {
            "perihelion_distance": "0.566072",
            "semi_major_axis": "",
            "eccentricity": "1.000077",
            "arg_perihelion": "243.6283",
            "long_asc_node": "108.0793",
            "inclination": "75.2398",
            "orbital_period": "",
            "perihelion_time": "2026-01-20 18:23:28",
            "epoch_time": "2025-10-04"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "7.0",
            "cobs_calc_n": "10",
            "cobs_obs_H": "7.6",
            "cobs_obs_n": "5.0",
            "astro_H": "7.0",
            "astro_n": "10",
            "mpc_H": "8.0",
            "mpc_n": "3.0"
        }
    }
}


# -------------------- 启动 --------------------
if __name__ == "__main__":
    root = tk.Tk()
    CometBrightnessApp(root)
    root.mainloop()
